import os
import pickle
import wave
from typing import List

import numpy as np
import opuslib_next
from pydub import AudioSegment


class OpusEncode:
    # 单例模式
    _instance = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self):
        # self.encoder = opuslib.Encoder(48000, 2, opuslib.APPLICATION_AUDIO)
        # self.encoder.set_bitrate(128000)
        # self.encoder.set_complexity(10)
        # self.encoder.set_inband_fec(True)
        # self.encoder.set_packet_loss_percentage(0)
        # self.encoder.set_signal(opuslib.SIGNAL_MUSIC)
        # self.encoder.set_vbr(True)
        # self.encoder.set_vbr_constraint(True)
        self.channel = 1
        self.sample_rate = 16000
        self.sample_width = 2
        self.opus_sample_rate = 16000
        self.opus_channel = 1
        self.opus_sample_width = 2
        self.opus_frame_time = 60 # ms,一帧时长60ms
        self.opus_frame_size = int(self.opus_sample_rate * self.opus_frame_time / 1000)

    def audio_to_opus(self, audio_path):
        # 获取文件类型
        file_type = os.path.splitext(audio_path)[1]
        # file_type = audio_path.split(".")[-1]
        if file_type:
            file_type = file_type.lstrip(".")

        # 读取音频文件,解析文件格式 MP3/wav，得到一个对象
        audio = AudioSegment.from_file(audio_path,format=file_type)

        # 设置参数
        audio = audio.set_channels(self.channel).set_frame_rate(self.sample_rate).set_sample_width(self.sample_width)

        # 获取音频时长 s
        # duration = audio.duration_seconds
        audio_duration  = len(audio) / 1000
        print("音频时长: %s s" % audio_duration )

        # 获取原始的PCM数据
        pcm_raw_data = audio.raw_data
        print("原始PCM数据长度: %s" % len(pcm_raw_data))

        # 通过opuslib进行编码,初始化一个编码器
        encoder = opuslib_next.Encoder(self.opus_sample_rate, self.opus_channel, opuslib_next.APPLICATION_AUDIO)

        # 获取每帧的采样数
        # frame_size = int(self.opus_sample_rate * self.opus_frame_time / 1000)
        frame_size = self.opus_frame_size

        # 获取每帧的采样字节数
        frame_size_bytes = frame_size * self.opus_channel * self.opus_sample_width

        opus_datas = []  # 存放编码后的数据
        # 优化，直接在外面进行补齐，不用每次循环重新判断
        if len(pcm_raw_data) % frame_size_bytes != 0:
            pcm_raw_data += b'\x00' * (frame_size_bytes - len(pcm_raw_data) % frame_size_bytes)
        # 通过编码器进行编码
        for i in range(0, len(pcm_raw_data), frame_size_bytes):  # 以每帧的字节数为步长，遍历pcm_raw_data
            # 获取一帧的采样数据
            frame_data = pcm_raw_data[i:i + frame_size_bytes]
            # # 补齐
            # if len(frame_data) < frame_size_bytes:
            #     frame_data += b'\x00' * (frame_size_bytes - len(frame_data))

            # 将字节数据转换为int16类型
            np_frame = np.frombuffer(frame_data, dtype=np.int16)

            # 将int16类型的数据转换为字节数
            np_frame = np_frame.tobytes()

            # 压缩编码,获取每一帧的opus数据
            encoded_frame = encoder.encode(np_frame, frame_size)

            # 添加到列表
            opus_datas.append(encoded_frame)

        return opus_datas, audio_duration

    # load 方法和save方法需要配对使用
    @staticmethod
    def save_opus_raw_data(opus_datas, output_file):
        with open(output_file, "wb") as pkl_file:
            pickle.dump(opus_datas, pkl_file) # type: ignore

    @staticmethod
    def load_opus_raw_data(input_path):
        with open(input_path, "rb") as f:
            opus_datas = pickle.load(f)
        return opus_datas

    @staticmethod
    def save_opus_custom_data(opus_datas, output_file):
        with open(output_file, "wb") as f:
            for frame_data in opus_datas:
                f.write(len(frame_data).to_bytes(4, byteorder="big"))
                f.write(frame_data)

    @staticmethod
    def load_opus_custom_data(input_path):
        print(f"[OpusEncode] load_opus_custom_data input_path: {input_path}")
        with open(input_path, "rb") as f:
            opus_datas = []
            while True:
                opus_data_len = int.from_bytes(f.read(4), byteorder="big")
                if opus_data_len == 0:
                    break
                frame_data = f.read(opus_data_len)
                opus_datas.append(frame_data)
            return opus_datas

    def opus_to_wav_file(self, output_file, opus_datas : List[bytes]) -> str:
        # output_file = "../asr/funasr/test.wav"
        print(f"[OpusEncode] opus_to_wav_file: {output_file}")
        # 判断文件是否存在，如果文件不存在则创建
        if not os.path.exists(output_file):
            # 创建父目录（自动递归创建多级目录）
            os.makedirs(os.path.dirname(output_file), exist_ok=True)
            # 创建空文件（可选，wave.open写入时会自动覆盖）
            # open(output_file, 'a').close()

        if not opus_datas:  # 新增空列表校验
            raise ValueError("输入的opus数据不能为空")
        # 通过opuslib进行解码,初始化一个解码器
        decoder = opuslib_next.Decoder(self.opus_sample_rate, self.opus_channel)
        pcm_raw_data = []
        for frame_data in opus_datas:
            try:
                pcm_frame = decoder.decode(frame_data, self.opus_frame_size)
                pcm_raw_data.append(pcm_frame)
            except opuslib_next.OpusError as e:
                print(f"解码失败: {e}")

        with wave.open(output_file, "wb") as wave_file:
            wave_file.setnchannels(self.channel)
            wave_file.setsampwidth(self.sample_width)
            wave_file.setframerate(self.sample_rate)
            print(f"PCM数据：{pcm_raw_data}")
            print(b"".join(pcm_raw_data))
            wave_file.writeframes(b"".join(pcm_raw_data))
            return output_file



if __name__ == "__main__":
    opus = OpusEncode()
    opus_data,duration = opus.audio_to_opus("../../test.mp3")
    print(opus_data)
    print(duration)
    opus.save_opus_raw_data(opus_data, "test.pkl")
    opus_data = opus.load_opus_raw_data("test.pkl")
    print(opus_data)
    opus.save_opus_custom_data(opus_data, "test.opus")
    opus_data = opus.load_opus_custom_data("test.opus")
    print(opus_data)

    opus_data1 = opus.opus_to_wav_file("../asr/funasr/test.wav",opus_data)
    print(opus_data1)